package com.hy.common.threadpool;

import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.concurrent.ThreadPoolExecutor;

/**
 * 自定义线程池-重写线程池的创建
 */
public class ThreadPoolCreate {
    private ThreadPoolConfigParam threadPoolConfigParam;

    public ThreadPoolCreate(){
        // 获取CPU核数 cpuNum
        int cpuNum = Runtime.getRuntime().availableProcessors();
        threadPoolConfigParam = new ThreadPoolConfigParam();
        threadPoolConfigParam.setThreadPoolCorePoolSize(cpuNum);
        threadPoolConfigParam.setThreadPoolMaxPoolSize(cpuNum);
        threadPoolConfigParam.setThreadPoolQueueCapacity(65535);
        threadPoolConfigParam.setWaitForTasksToCompleteOnShutdown(true);
        threadPoolConfigParam.setAwaitTerminationSeconds(60);
    }
    public ThreadPoolCreate(ThreadPoolConfigParam customThreadPoolConfig){
        if(customThreadPoolConfig == null){
            int cpuNum = Runtime.getRuntime().availableProcessors();
            threadPoolConfigParam = new ThreadPoolConfigParam();
            threadPoolConfigParam.setThreadPoolCorePoolSize(cpuNum);
            threadPoolConfigParam.setThreadPoolMaxPoolSize(cpuNum);
            threadPoolConfigParam.setThreadPoolQueueCapacity(65535);
            threadPoolConfigParam.setWaitForTasksToCompleteOnShutdown(true);
            threadPoolConfigParam.setAwaitTerminationSeconds(60);
        }else{
            this.threadPoolConfigParam = customThreadPoolConfig;
        }
    }

    public ThreadPoolTaskExecutor getTaskExecutorByType(ThreadPoolTypeEnum threadPoolTypeEnum) {
        int cpuNum = Runtime.getRuntime().availableProcessors();
        // CPU密集型（其中N为CPU的个数）建议：线程池大小 = cpuNum+1，建议：线程池大小 = cpuNum + 1
        if(ThreadPoolTypeEnum.CPU.getCode().equals(threadPoolTypeEnum.getCode())){
            return getTaskExecutor(cpuNum + 1);
        }
        // IO密集型（其中cpuNum为CPU的个数）建议：线程池大小 = 2N+1
        if(ThreadPoolTypeEnum.IO.getCode().equals(threadPoolTypeEnum.getCode())){
            return getTaskExecutor(2*cpuNum + 1);
        }
        return getTaskExecutor(cpuNum);
    }

    public ThreadPoolTaskExecutor getTaskExecutor(Integer nThreads) {
        ThreadPoolTaskExecutor executor = new ThreadPoolTaskExecutor();
        if (null == nThreads) {
            // 设置核心线程数
            executor.setCorePoolSize(threadPoolConfigParam.getThreadPoolCorePoolSize());
            // 设置最大线程数
            executor.setMaxPoolSize(threadPoolConfigParam.getThreadPoolMaxPoolSize());
        } else {
            executor.setCorePoolSize(nThreads);
            executor.setMaxPoolSize(nThreads);
        }
        // 设置队列容量
        executor.setQueueCapacity(threadPoolConfigParam.getThreadPoolQueueCapacity());
        // 设置拒绝策略
        executor.setRejectedExecutionHandler(new ThreadPoolExecutor.CallerRunsPolicy());
        // 等待所有的任务结束后再关闭线程池
        executor.setWaitForTasksToCompleteOnShutdown(threadPoolConfigParam.getWaitForTasksToCompleteOnShutdown());
        // 等待指定时间没关闭自动关闭
        executor.setAwaitTerminationSeconds(threadPoolConfigParam.getAwaitTerminationSeconds());
        executor.initialize();
        return executor;
    }



}
